package edu.berkeley.nlp.assignments;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

import edu.berkeley.nlp.langmodel.LanguageModel;
import edu.berkeley.nlp.util.Counter;
import edu.berkeley.nlp.util.CounterMap;

public class KneserNeyLanguageModel implements LanguageModel {

	static final String STOP = "</S>";

	
	double total = 0.0;
	// for unigram, c(w')
	Counter<String> wordCounter = new Counter<String>();

	double totalBigram = 0.0;
	// for bigram, c*(w,w')
    CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();
    CounterMap<String, String> bigramKNProb = new CounterMap<String, String>();

    // for trigram	key:w-2,w-1, value: w
	double totalTrigram = 0.0;
    CounterMap<String, String> trigramCounterMap = new CounterMap<String, String>();
    CounterMap<String, String> trigramKNProb = new CounterMap<String, String>();

    // for reverseTrigram	key: w-1, w, value: w-2
    CounterMap<String, String> revTrigramCounterMap = new CounterMap<String, String>();
    
    double D = 0.75;
    
	public List<String> generateSentence() {
		// TODO Auto-generated method stub
		return null;
	}

	public double getSentenceProbability(List<String> sentence) {
		// TODO Auto-generated method stub
		return 0;
	}


	public void KneserNeyTrigramSmoothing() {
		/*
		for (String wMinus2: wordCounter.keySet()) {
			for (String wMinus1: wordCounter.keySet()) {
				for (String w: wordCounter.keySet()) {
					double prob = 0.0;
					String wConcat = getConcatString(wMinus2, wMinus2);
					// higher order
					prob = Math.max(trigramCounterMap.get, 0.0);
					if (trigramCounterMap.containsKey(wConcat)) {
						
					} 
					
					
					// lower order
				
					//assign prob
					trigramKNProb.setCount(wConcat, w, prob);
				}
			}
		}*/
		for (String wConcat: trigramCounterMap.keySet()) {
			String wFirst = getFirstString(wConcat), wSecond = getSecondString(wConcat);
			for (String w: wordCounter.keySet()) {
				double prob = 0.0;
				// higher order
				prob = Math.max(trigramCounterMap.getCount(wConcat, w) - D, 0.0) / trigramCounterMap.getCounter(wConcat).totalCount();
			
				// lower order
				double probLower = 0.0;
				probLower = D / trigramCounterMap.getCounter(wConcat).totalCount();
				
				int typeCount = 0;
				for (String wLastTemp : bigramCounterMap.keySet()) {
					if (bigramCounterMap.getCount(wLastTemp, w) > 0)
						typeCount ++;
				}
				probLower *= typeCount;
				
				probLower *= revTrigramCounterMap.getCounter(getConcatString(wSecond, w)).keySet().size();
				double denominator = 0.0;
				for (String wTemp: wordCounter.keySet()) {
					denominator += revTrigramCounterMap.getCounter(getConcatString(wSecond, wTemp)).keySet().size();
				}
				probLower /= denominator;
				
				
				prob += probLower;
				//assign prob
				trigramKNProb.setCount(wConcat, w, prob);
			}
		}
	}
	
	public KneserNeyLanguageModel(Collection<List<String>> sentenceCollection, int ngram) {
		for (List<String> sentence : sentenceCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			stoppedSentence.add(STOP);
			String lastWord = null, lastLastWord = null;
			for (String word : stoppedSentence) {
				wordCounter.incrementCount(word, 1.0);
				
				if (ngram >= 2 && lastWord != null) {
					bigramCounterMap.incrementCount(lastWord, word, 1.0);
				}
				
				if (ngram == 3 && lastLastWord != null) {
					trigramCounterMap.incrementCount(getConcatString(lastLastWord, lastWord), word, 1.0);
					revTrigramCounterMap.incrementCount(getConcatString(lastWord, word), lastLastWord, 1.0);
				}
				lastLastWord = lastWord;
				lastWord = word;
			}
		}
		total = wordCounter.totalCount();
		totalBigram = bigramCounterMap.totalCount();	
		totalTrigram = trigramCounterMap.totalCount();
		
		if (ngram == 3)
			KneserNeyTrigramSmoothing();
	}
	
	

	//trigram key manipulation
	String CONCAT = "<concat>";
	public String getFirstString(String concatString) {
		for (int i = 0; i < concatString.length(); i++) {
			if (concatString.substring(0, i).endsWith(CONCAT)) 
				return concatString.substring(0, i - 1 - CONCAT.length());
		}
		return null;
	}
	

	public String getSecondString(String concatString) {
		for (int i = 0; i < concatString.length(); i++) {
			if (concatString.substring(0, i).endsWith(CONCAT)) 
				return concatString.substring(i);
		}
		return null;
	}
	
	public String getConcatString(String first, String second) {
		String concat = "";
		concat = first + CONCAT + second;
		return concat;
	}
}
